import logging
import os
from typing import List
from sentence_transformers import SentenceTransformer
import faiss
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

"""FAISS 比 ChromaDB 更轻量，适合纯 CPU 环境（需安装 faiss-cpu 或 faiss-gpu）
    FAISS 仅存储向量，需自行维护原始文档列表（代码中的 self.documents）
    大数据集（>1M条）建议使用 IndexIVFFlat 或 IndexHNSWFlat
    如需恢复完整功能（如元数据过滤），可结合 SQLite 使用
    混合搜索（提高召回率）：self.index = faiss.IndexHNSWFlat(self.embedding_dim, 32)
"""

class SimpleRAG:
    def __init__(self, embedding_model_path: str = None, faiss_index_path: str = "./faiss_index"):
        """
        初始化RAG系统
        :param embedding_model_path: 本地嵌入模型路径（可选）
        :param faiss_index_path: FAISS索引保存路径
        """
        # 初始化嵌入模型
        if embedding_model_path and os.path.exists(embedding_model_path):
            self.embedding_model = SentenceTransformer(embedding_model_path)
        else:
            self.embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

        # 初始化FAISS
        self.faiss_index_path = faiss_index_path
        self.documents = []  # 存储原始文档
        self.embedding_dim = self.embedding_model.get_sentence_embedding_dimension()

        # 加载或创建FAISS索引
        if os.path.exists(faiss_index_path):
            self.index = faiss.read_index(faiss_index_path)
            logging.info(f"从 {faiss_index_path} 加载已有FAISS索引")
        else:
            self.index = faiss.IndexFlatL2(self.embedding_dim)
            logging.info("新建FAISS索引")

        # 初始化ChatLM-mini-Chinese模型
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        self.tokenizer = AutoTokenizer.from_pretrained("D:/ideaSpace/MyPython/models/ChatLM-mini-Chinese")
        self.llm_model = AutoModelForSeq2SeqLM.from_pretrained("D:/ideaSpace/MyPython/models/ChatLM-mini-Chinese").to(self.device)

    def add_documents(self, documents: List[str]):
        """
        添加文档到知识库
        :param documents: 文档列表
        """
        # 生成嵌入向量
        embeddings = self.embedding_model.encode(documents, convert_to_numpy=True)

        # 添加到FAISS索引
        if self.index.ntotal == 0:
            self.index.add(embeddings)
        else:
            self.index.add(embeddings)

        # 保存原始文档
        self.documents.extend(documents)

        # 持久化索引
        faiss.write_index(self.index, self.faiss_index_path)
        logging.info(f"成功添加 {len(documents)} 个文档到FAISS索引")

    def query(self, question: str, top_k: int = 3) -> str:
        """
        查询RAG系统
        :param question: 用户问题
        :param top_k: 返回最相关的文档数量
        :return: LLM生成的回答
        """
        # 1. 检索相关文档
        query_embedding = self.embedding_model.encode([question], convert_to_numpy=True)
        distances, indices = self.index.search(query_embedding, top_k)

        # 获取相关文档
        relevant_docs = [self.documents[i] for i in range(len(indices[0]))]
        context = "\n\n".join(relevant_docs)

        # 2. 构建提示词
        prompt = f"##提问:\n基于以下上下文回答：\n{context}\n问题：{question}\n##回答:\n"

        # 3. 调用模型生成回答
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
        outputs = self.llm_model.generate(
            input_ids,
            max_new_tokens=256,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.tokenizer.pad_token_id
        )
        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

        # 提取回答部分
        return answer.split("##回答:")[-1].strip()

# 使用示例
if __name__ == "__main__":
    # 初始化RAG（默认从HuggingFace下载嵌入模型）
    rag = SimpleRAG("D:/ideaSpace/MyPython/models/bge-small-zh-v1.5")

    # 添加文档
    documents = [
        "人工智能(AI)是计算机科学的一个分支，致力于开发能模拟人类智能的系统。",
        "机器学习是AI的核心技术之一，它使计算机能够从数据中学习和改进。",
        "深度学习是机器学习的一个子领域，使用多层神经网络处理复杂问题。"
    ]
    rag.add_documents(documents)

    # 查询示例
    question = "机器学习和深度学习有什么区别？"
    answer = rag.query(question)
    print(f"问题: {question}")
    print(f"回答: {answer}")